﻿#region [ License information          ]
/* ************************************************************
 *
 * Copyright (c) Bjørn Bouet Smith, 2010, 2011
 *
 * This source code is subject to terms and conditions of 
 * Microsoft Public License (Ms-PL).
 * 
 * A copy of the license can be found in the license.txt
 * file at the root of this distribution. If you can not 
 * locate the License, please send an email to bjornsmith@gmail.com
 * 
 * By using this source code in any fashion, you are 
 * agreeing to be bound by the terms of the Microsoft 
 * Permissive License.
 *
 * You must not remove this notice, or any other, from this
 * software.
 *
 * ************************************************************/
#endregion

using System;
using System.Collections.Generic;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Text;
using Smith.AMC.Operations;
using Wintellect.Threading.AsyncProgModel;

namespace Smith.AMC.Network
{
    public class UdpClientSocket : ClientSocket
    {
        /// <summary>
        /// Initializes a new instance of the <see cref="UdpClientSocket"/> class.
        /// </summary>
        /// <param name="ipEndPoint">The ip end point.</param>
        /// <param name="socketType">Type of the socket.</param>
        /// <param name="protocol">The protocol.</param>
        public UdpClientSocket(
            IPEndPoint ipEndPoint,
            SocketType socketType,
            ProtocolType protocol)
            : base(ipEndPoint, socketType, protocol)
        {
        }

        /// <summary>
        /// Method that does the actual write.
        /// </summary>
        /// <typeparam name="T">Type of data the operation contains</typeparam>
        /// <param name="ae">The async enumerator.</param>
        /// <param name="operation">The operation.</param>
        /// <returns>
        /// An Ienumerator that the async enumenrator iterates over
        /// </returns>
        protected override IEnumerator<int> DoWrite<T>(AsyncEnumerator<OperationResponse<T>> ae, Operation operation)
        {
            if (!operation.Name.Equals(OperationNames.Get))
            {
                throw new InvalidOperationException("Only GetOperations are allowed to use UDP sockets");
            }

            if (!Connected)
            {
                BeginConnect(IpEndPoint, ae.End(), null);
                yield return 1;
                EndConnect(ae.DequeueAsyncResult());
            }

            string headerLine = operation.ToString();

            byte[] command = Encoding.ASCII.GetBytes(headerLine);
            byte[] frameHeader = CreateFrameHeader(1, 0, 1);

            ArraySegment<byte>[] arraySegments = new []
            {
                new ArraySegment<byte>(frameHeader),
                new ArraySegment<byte>(command)
            };

            SocketError err;

            // Send buffer
            BeginSend(arraySegments, SocketFlags.None, out err, ae.End(), null);

            yield return 1;
            EndSend(ae.DequeueAsyncResult());

            // Memcached server sends at most 1400 bytes via UDP
            const int MaxBytesToRead = 1400;
            byte[] readBuffer = new byte[MaxBytesToRead];
            BeginReceive(readBuffer, 0, MaxBytesToRead, SocketFlags.None, out err, ae.End(), null);
            yield return 1;
            int bytesRead = EndReceive(ae.DequeueAsyncResult());

            byte[][] destinationArray = null;
            DecodeServerResponsePackage(readBuffer, bytesRead, ref destinationArray);
            if (destinationArray.Length > 1)
            {
                // More data to expect
                BeginReceive(readBuffer, 0, MaxBytesToRead, SocketFlags.None, out err, ae.End(), null);
                yield return 1;
                bytesRead = EndReceive(ae.DequeueAsyncResult());
                DecodeServerResponsePackage(readBuffer, bytesRead, ref destinationArray);
            }
        }

        private static byte[] CreateFrameHeader(short requestID, short sequenceNumber, short numberOfDataGrams)
        {
            byte[] requestIDs = BitConverter.GetBytes(IPAddress.HostToNetworkOrder(requestID));
            byte[] sequenceNumbers = BitConverter.GetBytes(IPAddress.HostToNetworkOrder(sequenceNumber));

            byte[] numberOfDataGramss = BitConverter.GetBytes(IPAddress.HostToNetworkOrder(numberOfDataGrams));

            byte[] reserved = BitConverter.GetBytes(IPAddress.HostToNetworkOrder((short)0));
            using (MemoryStream ms = new MemoryStream())
            {
                ms.Write(requestIDs, 0, requestIDs.Length);
                ms.Write(sequenceNumbers, 0, sequenceNumbers.Length);
                ms.Write(numberOfDataGramss, 0, numberOfDataGramss.Length);
                ms.Write(reserved, 0, reserved.Length);
                return ms.ToArray();
            }
        }

        /// <summary>
        /// Decodes the server response package and put the data into the destination array, which is a multidimensional array
        /// where first dimension is the package number, and the second dimension is the actual data
        /// 
        /// i.e. destinationArray[0]= first package's byte array
        /// destinationArray[1]= first package's byte array
        /// etc.
        /// </summary>
        /// <param name="serverPackage">The server package.</param>
        /// <param name="dataLength">Length of the data.</param>
        /// <param name="destinationArray">The destination array.</param>
        private static void DecodeServerResponsePackage(byte[] serverPackage, int dataLength, ref byte[][] destinationArray)
        {
            const int FrameHeaderLength = 8;
            if (serverPackage.Length < 13)
            {
                // Something went wrong, frame header is 8 bytes
                // And the smallest response is END\r\n which is 5 bytes
                throw new NetworkException(string.Format("UDP package contains invalid size:{0} - 13 bytes is the smallest valid package size", serverPackage.Length));
            }
            
            KeyValuePair<short, short> header = DecodeFrameHeader(serverPackage);
            short totalNumberOfPackaged = header.Value;
            short packageNumber = header.Key;

            if (destinationArray == null)
            {
                destinationArray = new byte[totalNumberOfPackaged][];

            }

            const int StartPosition = FrameHeaderLength;
            int readLength = dataLength - FrameHeaderLength;

            destinationArray[packageNumber] = new byte[readLength];
            
            // Copy valid bytes to the destination array
            Array.Copy(serverPackage, StartPosition, destinationArray[packageNumber], 0, readLength);
           
        }

        /// <summary>
        /// Decodes the frame header.
        /// Frame header is 8 bytes
        /// 0-1 Request ID
        /// 2-3 Sequence number
        /// 4-5 Total number of datagrams in this message
        /// 6-7 Reserved for future use; must be 0
        /// </summary>
        /// <param name="serverPackage">The server package.</param>
        /// <returns>A key value pair, with the current package number out of the total packages.
        /// Key is current package number
        /// Value is total number of packages</returns>
        private static KeyValuePair<short, short> DecodeFrameHeader(byte[] serverPackage)
        {
            // Frame header is 8 bytes
            // 0-1 Request ID
            // 2-3 Sequence number
            // 4-5 Total number of datagrams in this message
            // 6-7 Reserved for future use; must be 0

            short sequenceNumber = IPAddress.NetworkToHostOrder(BitConverter.ToInt16(serverPackage, 2));
            short totalNumberOfPackages = IPAddress.NetworkToHostOrder(BitConverter.ToInt16(serverPackage, 4));
            return new KeyValuePair<short, short>(sequenceNumber, totalNumberOfPackages);
        }

    }
}
